package com.cheems.cheemsspringaialibaba.Graph.config;

import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.OverAllStateFactory;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncEdgeAction;
import com.alibaba.cloud.ai.graph.action.EdgeAction;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.memory.jdbc.MysqlChatMemoryRepository;
import com.cheems.cheemsspringaialibaba.Graph.node.*;
import com.cheems.cheemsspringaialibaba.Tools.SearchFromDify;
import com.cheems.cheemsspringaialibaba.Tools.SearchStudentSchedule;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 * @author JTB
 */
@Configuration
public class HenauRoutingGraphConfiguration {

    //可以选择的问题路由模式
    public static Map<String, Object> availableRoutes = new HashMap<>();

    static {
        availableRoutes.put("searchExam", "你是一个成绩查询小助手,你只可以帮助同学进行成绩查询和课程查询");
        availableRoutes.put("rag", "你是一个河南农业大学校内知识专家,可以对农大校内的知识进行回答");
        availableRoutes.put("normal","你是一个普通的小助手,你可以处理其他角色无法处理的事情,查询天气 回答学科问题 编程等");
    }

    //创建路由图 进行流程编排
    @Bean
    public StateGraph HenauRoutingGraph(ChatModel chatModel, JdbcTemplate jdbcTemplate) throws GraphStateException {
        // 构造 ChatMemoryRepository 和 ChatMemory
        ChatMemoryRepository chatMemoryRepository = MysqlChatMemoryRepository.mysqlBuilder()
                .jdbcTemplate(jdbcTemplate)
                .build();
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(chatMemoryRepository)
                .build();
        ChatClient client = ChatClient.builder(chatModel).defaultAdvisors(new SimpleLoggerAdvisor() , MessageChatMemoryAdvisor.builder(chatMemory).build()).build();

        //全局变量管理工厂
        OverAllStateFactory factory = () -> {
            OverAllState s = new OverAllState();
            s.registerKeyAndStrategy("inputText", new ReplaceStrategy());
            s.registerKeyAndStrategy("selectionLlm", new ReplaceStrategy());
            s.registerKeyAndStrategy("result" , new ReplaceStrategy());
            return s;
        };

        //创建路由图
        StateGraph graph = new StateGraph("HenauRoutingGraph", factory.create())
                .addNode("routing", node_async(new LlmRoutingNode.LlmCallRouterNode(client, "inputText")))
                .addNode("rag", node_async(new RagNode(client)))
                //.addNode("normal", node_async(new NormalNode(client)))
                .addNode("searchExam", node_async(new SearchExamNode(client)))
                // 构建并行边：使用单条边携带多目标
                .addEdge(START, "routing")
                //添加条件边，根据路由节点的输出结果决定下一步流向：
                .addConditionalEdges("routing", AsyncEdgeAction.edge_async(new EdgeAction() {
                    @Override
                    public String apply(OverAllState state) {
                        String selection = (String) state.value("selectionLlm").orElse("");
                        return selection;
                    }
                }), Map.of("rag", "rag",  "searchExam", "searchExam"))
                .addEdge("rag", END)
                .addEdge("searchExam", END);
                //.addEdge("normal", END);

        // 可视化
        GraphRepresentation representation = graph.getGraph(GraphRepresentation.Type.PLANTUML,
                "Routing  flow");
        System.out.println("\n=== Routing  UML Flow ===");
        System.out.println(representation.content());
        System.out.println("==================================\n");

        return graph;
    }

}
